package ai

import (
	"context"
	"fmt"
	"log"
	"sync"
	"time"
	"xagent/internal/memory"
	protov1 "xagent/proto/v1"

	"github.com/google/uuid"
)

// AdaptationMetrics tracks adaptation performance
type AdaptationMetrics struct {
	AdjustmentsMade   int
	AdaptationSuccess float64
	ResponseTime      time.Duration
	LastUpdate        time.Time
}

// AdaptationManager handles AI adaptation capabilities
type AdaptationManager struct {
	memoryManager     memory.MemoryManager
	metrics           *AdaptationMetrics
	strategies        map[string]*AdaptationStrategy
	activeAdaptations map[string]*Adaptation
	mutex             sync.RWMutex
	updateInterval    time.Duration
}

// AdaptationStrategy represents a strategy for handling specific situations
type AdaptationStrategy struct {
	Name        string
	Triggers    []string
	Actions     []AdaptiveAction
	SuccessRate float64
	UsageCount  int
	LastUsed    time.Time
}

// AdaptiveAction represents an action to take in response to a situation
type AdaptiveAction struct {
	Type       string
	Parameters map[string]interface{}
	Priority   int
}

// Adaptation represents an ongoing adaptation process
type Adaptation struct {
	Strategy    *AdaptationStrategy
	StartTime   time.Time
	Status      string
	Progress    float64
	Adjustments []Adjustment
}

// Adjustment represents a single adaptation adjustment
type Adjustment struct {
	Type      string
	Value     interface{}
	Timestamp time.Time
	Success   bool
}

// NewAdaptationManager creates a new adaptation manager
func NewAdaptationManager(memoryManager memory.MemoryManager) *AdaptationManager {
	return &AdaptationManager{
		memoryManager:     memoryManager,
		metrics:           &AdaptationMetrics{},
		strategies:        make(map[string]*AdaptationStrategy),
		activeAdaptations: make(map[string]*Adaptation),
		updateInterval:    time.Minute * 5,
	}
}

// ProcessFeedback processes feedback and adapts behavior
func (am *AdaptationManager) ProcessFeedback(ctx context.Context, feedback *protov1.FeedbackMessage) error {
	// Analyze feedback
	strategy := am.selectStrategy(feedback)
	if strategy == nil {
		strategy = am.createNewStrategy(feedback)
	}

	// Create adaptation
	adaptation := &Adaptation{
		Strategy:  strategy,
		StartTime: time.Now(),
		Status:    "active",
	}

	// Start adaptation process
	am.mutex.Lock()
	am.activeAdaptations[feedback.Id] = adaptation
	am.mutex.Unlock()

	// Apply adaptive actions
	for _, action := range strategy.Actions {
		adjustment := am.applyAction(ctx, action, feedback)
		adaptation.Adjustments = append(adaptation.Adjustments, adjustment)

		if !adjustment.Success {
			adaptation.Status = "failed"
			break
		}

		adaptation.Progress += 1.0 / float64(len(strategy.Actions))
	}

	// Update strategy success rate
	am.updateStrategyMetrics(strategy, adaptation.Status == "active")

	return nil
}

// OptimizeResources dynamically optimizes resource allocation
func (am *AdaptationManager) OptimizeResources(ctx context.Context, resources map[string]float64) (map[string]float64, error) {
	optimized := make(map[string]float64)
	for resource := range resources {
		// Get relevant memories for this resource
		memories, err := am.memoryManager.GetRelevantMemories(ctx, resource, 10, 0.7)
		if err != nil {
			return nil, fmt.Errorf("failed to get memories for resource %s: %v", resource, err)
		}

		// Calculate optimal allocation based on historical data
		optimized[resource] = am.calculateOptimalAllocation(memories)
	}

	// Record optimization for future reference
	am.recordOptimization(optimized)
	return optimized, nil
}

// HandleError adapts error handling strategy based on past experiences
func (am *AdaptationManager) HandleError(ctx context.Context, err error, task *protov1.Task) error {
	// Get relevant error handling memories
	memories, err := am.memoryManager.GetRelevantMemories(ctx, err.Error(), 5, 0.8)
	if err != nil {
		return fmt.Errorf("failed to get error handling memories: %v", err)
	}

	strategy := am.selectErrorStrategy(memories)
	if strategy == nil {
		strategy = am.createErrorStrategy(err)
	}

	// Apply error handling actions
	for _, action := range strategy.Actions {
		if err := am.applyErrorAction(ctx, action, task); err != nil {
			return fmt.Errorf("failed to apply error action: %v", err)
		}
	}

	return nil
}

// StartAdaptationLoop starts periodic adaptation updates
func (am *AdaptationManager) StartAdaptationLoop(ctx context.Context) {
	go func() {
		ticker := time.NewTicker(am.updateInterval)
		defer ticker.Stop()

		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				am.updateAdaptations(ctx)
			}
		}
	}()
}

// Helper functions

func (am *AdaptationManager) selectStrategy(feedback *protov1.FeedbackMessage) *AdaptationStrategy {
	am.mutex.RLock()
	defer am.mutex.RUnlock()

	var bestStrategy *AdaptationStrategy
	bestScore := -1.0

	for _, strategy := range am.strategies {
		score := am.calculateStrategyScore(strategy, feedback)
		if score > bestScore {
			bestStrategy = strategy
			bestScore = score
		}
	}

	return bestStrategy
}

func (am *AdaptationManager) createNewStrategy(feedback *protov1.FeedbackMessage) *AdaptationStrategy {
	feedbackType := feedback.Metadata["type"]
	strategy := &AdaptationStrategy{
		Name:     fmt.Sprintf("strategy_%d", len(am.strategies)),
		Triggers: []string{feedbackType},
		Actions:  am.generateActions(feedback),
	}

	am.mutex.Lock()
	am.strategies[strategy.Name] = strategy
	am.mutex.Unlock()

	return strategy
}

func (am *AdaptationManager) calculateStrategyScore(strategy *AdaptationStrategy, feedback *protov1.FeedbackMessage) float64 {
	if strategy.UsageCount == 0 {
		return 0
	}

	// Base score on success rate
	score := strategy.SuccessRate

	// Adjust for recency
	timeSinceUse := time.Since(strategy.LastUsed)
	recencyFactor := 1.0 / (1.0 + float64(timeSinceUse.Hours()))
	score *= recencyFactor

	// Check trigger match
	feedbackType := feedback.Metadata["type"]
	for _, trigger := range strategy.Triggers {
		if trigger == feedbackType {
			score *= 1.5 // Boost score for matching triggers
			break
		}
	}

	return score
}

func (am *AdaptationManager) generateActions(feedback *protov1.FeedbackMessage) []AdaptiveAction {
	var actions []AdaptiveAction

	feedbackType := feedback.Metadata["type"]
	switch feedbackType {
	case "performance":
		actions = append(actions, AdaptiveAction{
			Type: "adjust_resources",
			Parameters: map[string]interface{}{
				"factor": 1.2,
			},
			Priority: 1,
		})
	case "error":
		actions = append(actions, AdaptiveAction{
			Type: "increase_retries",
			Parameters: map[string]interface{}{
				"max_retries": 3,
				"delay":       time.Second * 5,
			},
			Priority: 2,
		})
	case "quality":
		actions = append(actions, AdaptiveAction{
			Type: "enhance_validation",
			Parameters: map[string]interface{}{
				"strict_mode": true,
			},
			Priority: 1,
		})
	}

	return actions
}

func (am *AdaptationManager) applyAction(ctx context.Context, action AdaptiveAction, feedback *protov1.FeedbackMessage) Adjustment {
	adjustment := Adjustment{
		Type:      action.Type,
		Timestamp: time.Now(),
	}

	// Apply action based on type
	switch action.Type {
	case "adjust_resources":
		factor := action.Parameters["factor"].(float64)
		adjustment.Value = factor
		adjustment.Success = true
	case "increase_retries":
		maxRetries := action.Parameters["max_retries"].(int)
		adjustment.Value = maxRetries
		adjustment.Success = true
	case "enhance_validation":
		strictMode := action.Parameters["strict_mode"].(bool)
		adjustment.Value = strictMode
		adjustment.Success = true
	default:
		adjustment.Success = false
	}

	return adjustment
}

func (am *AdaptationManager) calculateOptimalAllocation(memories []*memory.MemoryEntry) float64 {
	if len(memories) == 0 {
		return 1.0
	}

	var totalValue, totalWeight float64
	for _, mem := range memories {
		if value, ok := mem.Metadata[string(mem.Type)].(float64); ok {
			// Weight recent memories more heavily
			timeSince := time.Since(mem.UpdatedAt)
			weight := 1.0 / (1.0 + float64(timeSince.Hours()))

			totalValue += value * weight
			totalWeight += weight
		}
	}

	if totalWeight == 0 {
		return 1.0
	}

	return totalValue / totalWeight
}

func (am *AdaptationManager) recordOptimization(optimized map[string]float64) {
	memory := &memory.MemoryEntry{
		ID:      uuid.NewString(),
		Type:    memory.MemoryTypeResult,
		Content: "Resource optimization result",
		Metadata: map[string]interface{}{
			"optimized": optimized,
			"success":   true,
		},
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
	}

	// Store optimization result in memory
	if err := am.memoryManager.AddToShortTerm(context.Background(), memory); err != nil {
		// Log error but don't fail the operation
		log.Printf("Failed to record optimization: %v", err)
	}
}

func (am *AdaptationManager) selectErrorStrategy(memories []*memory.MemoryEntry) *AdaptationStrategy {
	if len(memories) == 0 {
		return nil
	}

	strategies := make(map[string]int)
	for _, mem := range memories {
		if strategy, ok := mem.Metadata["strategy"].(string); ok {
			if mem.Metadata["success"].(bool) {
				strategies[strategy]++
			}
		}
	}

	var bestStrategy string
	maxSuccess := 0
	for strategy, successes := range strategies {
		if successes > maxSuccess {
			bestStrategy = strategy
			maxSuccess = successes
		}
	}

	return am.strategies[bestStrategy]
}

func (am *AdaptationManager) createErrorStrategy(err error) *AdaptationStrategy {
	return &AdaptationStrategy{
		Name:     fmt.Sprintf("error_strategy_%d", len(am.strategies)),
		Triggers: []string{"error", err.Error()},
		Actions: []AdaptiveAction{
			{
				Type: "retry",
				Parameters: map[string]interface{}{
					"max_attempts": 3,
					"delay":        time.Second * 5,
				},
				Priority: 1,
			},
			{
				Type: "log_error",
				Parameters: map[string]interface{}{
					"level": "error",
				},
				Priority: 2,
			},
		},
	}
}

func (am *AdaptationManager) applyErrorAction(ctx context.Context, action AdaptiveAction, task *protov1.Task) error {
	switch action.Type {
	case "retry":
		maxAttempts := action.Parameters["max_attempts"].(int)
		delay := action.Parameters["delay"].(time.Duration)

		for attempt := 1; attempt <= maxAttempts; attempt++ {
			// Retry task execution
			if err := am.retryTask(ctx, task); err == nil {
				return nil
			}
			time.Sleep(delay)
		}
		return fmt.Errorf("max retry attempts reached")

	case "log_error":
		// Log error details
		return nil

	default:
		return fmt.Errorf("unknown error action type: %s", action.Type)
	}
}

func (am *AdaptationManager) retryTask(ctx context.Context, task *protov1.Task) error {
	// Implement task retry logic
	return nil
}

func (am *AdaptationManager) updateAdaptations(ctx context.Context) {
	am.mutex.Lock()
	defer am.mutex.Unlock()

	now := time.Now()

	// Update active adaptations
	for id, adaptation := range am.activeAdaptations {
		if now.Sub(adaptation.StartTime) > time.Hour {
			// Complete long-running adaptations
			adaptation.Status = "completed"
			delete(am.activeAdaptations, id)
		}
	}

	// Update metrics
	am.metrics.LastUpdate = now
}

func (am *AdaptationManager) updateStrategyMetrics(strategy *AdaptationStrategy, success bool) {
	strategy.UsageCount++
	strategy.LastUsed = time.Now()

	if success {
		// Update success rate with exponential moving average
		strategy.SuccessRate = 0.9*strategy.SuccessRate + 0.1
	} else {
		strategy.SuccessRate = 0.9 * strategy.SuccessRate
	}
}

// GetMetrics returns current adaptation metrics
func (am *AdaptationManager) GetMetrics() *AdaptationMetrics {
	am.mutex.RLock()
	defer am.mutex.RUnlock()
	return am.metrics
}
